import json

from langchain.agents import create_react_agent
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.messages import SystemMessage
from langchain_core.prompts import ChatPromptTemplate


def get_current_weather(location, unit="fahrenheit"):
    """Get the current weather in a given location"""
    if "tokyo" in location.lower():
        return json.dumps({"location": "Tokyo", "temperature": "10", "unit": unit})
    elif "san francisco" in location.lower():
        return json.dumps({"location": "San Francisco", "temperature": "72", "unit": unit})
    elif "paris" in location.lower():
        return json.dumps({"location": "Paris", "temperature": "22", "unit": unit})
    else:
        return json.dumps({"location": location, "temperature": "unknown"})



print(get_current_weather('tokyo'))
# 初始化OpenAI LLM模型
llm = ChatZhipuAI(model="GLM-4-Flash")

tools = [
    {
        "name": "get_current_weather",
        "description": "Get the current weather in a given location",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The city and state, e.g. San Francisco, CA",
                },
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
            },
            "required": ["location"],
        },
    }
]
_EXTRACTION_TEMPLATE = """
Passage:
{input}
"""
SYSTEM_PROMPT_TEMPLATE = """
Extra the passage as best as you can. You have access to the following tools, please return the tool name:
{tools}
"""
question_prompt = ChatPromptTemplate.from_role_strings(
    [["system", SYSTEM_PROMPT_TEMPLATE], ["user", _EXTRACTION_TEMPLATE]])  ## sequence@!!
messages = [
    question_prompt[0].format(tools=tools),
    question_prompt[1].format(input="What’s the weather like in Boston right now?")
]


result = llm(messages)

print(result)

# system_message = SystemMessage(content='现在去 tokyo 旅游如何')
# agent_executor = create_react_agent(llm, tools, state_modifier=system_message, debug=True)

